# 模块导入
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Flatten

# 导入数据集，分别为输入特征和标签
mnist = tf.keras.datasets.mnist
# (x_train, y_train)：（训练集输入特征，训练集标签)
# (x_test, y_test)：(测试集输入特征，测试集标签)
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 对输入网络的输入特征进行归一化，使原本0到255之间的灰度值，变为0到1之间的数值
# （把输入特征的数值变小更适合神经网络吸收）
x_train, x_test = x_train / 255.0, x_test / 255.0


# 搭建网络结构
class MnistModel(Model):
    def __init__(self):
        super(MnistModel, self).__init__()
        self.flatten = Flatten()  # 将输入特征拉直为一维数组，也就是拉直为28*28=784个数值
        self.d1 = Dense(128, activation="relu")  # 第一层网络128个神经元，使用relu激活函数
        self.d2 = Dense(10, activation="softmax")  # 第二层网络10个神经元，使用softmax激活函数，使输出符合概率分布

    def call(self, x):
        x = self.flatten(x)
        x = self.d1(x)
        y = self.d2(x)
        return y


model = MnistModel()

# 配置训练方法
model.compile(optimizer="adam",  # 优化器
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),  # 损失函数，输出是概率分布，from_logits=False
              metrics=["sparse_categorical_accuracy"])  # 数据集中的标签是数值，神经网络输出y是概率分布

# 执行训练过程
model.fit(x_train,  # 训练集输入特征
          y_train,  # 训练集标签
          batch_size=32,  # 每次喂入网络32组数据
          epochs=5,  # 数据集迭代5次
          validation_data=(x_test, y_test),  # 测试集输入特征，测试集标签
          validation_freq=1)  # 每迭代1次训练集执行一次测试集的评测

# 打印出网络结构和参数统计
model.summary()